# 批量 将obj 转换为ply 格式

import numpy as np
from PIL import Image
import os

def parse_obj_with_texture(obj_path):
    vertices = []           # 存储顶点坐标 [v1, v2, ...]
    tex_coords = []         # 存储纹理坐标 [vt1, vt2, ...]
    vertex_indices = []     # 面的顶点索引
    tex_indices = []        # 面的纹理坐标索引
    
    with open(obj_path, 'r') as f:
        for line in f:
            if line.startswith('v '):  # 顶点坐标
                parts = line.strip().split()
                vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])
            elif line.startswith('vt '):  # 纹理坐标
                parts = line.strip().split()
                tex_coords.append([float(parts[1]), float(parts[2])])
            elif line.startswith('f '):  # 面定义（格式：f v1/vt1 v2/vt2 ...）
                parts = line.strip().split()[1:]
                for part in parts:
                    v_idx, vt_idx, _ = (part + '//').split('/')[:3]  # 处理可能的缺失值
                    vertex_indices.append(int(v_idx) - 1)  # OBJ索引从1开始
                    tex_indices.append(int(vt_idx) - 1) if vt_idx else None
    
    return np.array(vertices), np.array(tex_coords), vertex_indices, tex_indices


def load_texture_from_mtl(mtl_path, texture_path_file):
    with open(mtl_path, 'r') as f:
        for line in f:
            if line.startswith('map_Kd'):
                texture_path = line.split()[1].strip()
                break
    if texture_path:
        return np.array(Image.open(os.path.join(texture_path_file, texture_path)))  # 返回纹理图像（H,W,3）
    return None
 

def texcoord_to_rgb(tex_coords, texture):
    # 确保纹理坐标在 [0,1] 范围内（OBJ允许超出）
    tex_coords = np.clip(tex_coords, 0, 1)
    # 转换为纹理图像的像素坐标
    u = (tex_coords[:, 0] * (texture.shape[1] - 1)).astype(int)
    v = ((1 - tex_coords[:, 1]) * (texture.shape[0] - 1)).astype(int)  # 纹理V轴通常反向
    # 采样RGB颜色
    return texture[v, u]  # 返回 (N,3) 的RGB数组


def save_ply_with_color(vertices, colors, ply_path):
    with open(ply_path, 'w') as f:
        f.write("ply\n")
        f.write("format ascii 1.0\n")
        f.write(f"element vertex {len(vertices)}\n")
        f.write("property float x\n")
        f.write("property float y\n")
        f.write("property float z\n")
        f.write("property uchar red\n")
        f.write("property uchar green\n")
        f.write("property uchar blue\n")
        f.write("end_header\n")
        for v, c in zip(vertices, colors):
            f.write(f"{v[0]} {v[1]} {v[2]} {int(c[0])} {int(c[1])} {int(c[2])}\n")


def save_ply(obj_path, mtl_path, texture_path, output_npy):
    #  obj mtl 路径
    # obj_path = "E:\\data\\3DpointCloud\\松霖人脸数据\\face_clear\\1006_1381912792830971904-17495152859693Dresult\\3Dface_result.obj"
    # mtl_path = "E:\\data\\3DpointCloud\\松霖人脸数据\\face_clear\\1006_1381912792830971904-17495152859693Dresult\\3Dface_result.mtl"
    # texture_path = "E:\\data\\3DpointCloud\\松霖人脸数据\\face_clear\\1006_1381912792830971904-17495152859693Dresult"

    vertices, tex_coords, v_indices, vt_indices = parse_obj_with_texture(obj_path)
    texture = load_texture_from_mtl(mtl_path, texture_path)  # 需确保.mtl和纹理图片路径正确

    # 确保纹理坐标在[0,1]范围内
    tex_coords = np.clip(tex_coords, 0, 1)
    # 初始化顶点颜色
    vertex_colors = np.zeros((len(vertices), 3))

    if texture is not None:
        # 预计算纹理尺寸边界
        max_u = texture.shape[1] - 1
        max_v = texture.shape[0] - 1
        
        for i in range(len(v_indices)):
            if vt_indices[i] is not None:
                # 安全计算像素坐标
                u = tex_coords[vt_indices[i], 0]
                v = tex_coords[vt_indices[i], 1]
                
                u_pixel = int(u * max_u)
                v_pixel = int((1 - v) * max_v)  # 翻转V轴
                
                # 二次边界检查
                u_pixel = min(max_u, max(0, u_pixel))
                v_pixel = min(max_v, max(0, v_pixel))
                
                vertex_colors[v_indices[i]] = texture[v_pixel, u_pixel]
    else:
        vertex_colors = np.ones((len(vertices), 3)) * 255  # 白色

    save_ply_with_color(vertices, vertex_colors, output_npy)


#  找到循环列表

data_path = "E:\\data\\3DpointCloud\\松霖人脸数据\\face_clear" 

file_list = os.listdir(data_path)
# print(file_list[:10])

for i in file_list:
    face_file_list = os.listdir(os.path.join(data_path, i))
    obj_file_list = [f for f in face_file_list if f.endswith(".obj")]
    for obj in obj_file_list:
        file_name = obj[:-4]
        # print(file_name)
        os.makedirs(os.path.join("E:\data\\3DpointCloud\\松霖人脸数据\\face_seg", i+file_name), exist_ok=True)
        save_ply(os.path.join(data_path, i, obj), 
                 os.path.join(data_path, i, file_name+".mtl"),
                 os.path.join(data_path, i),
                 os.path.join("E:\data\\3DpointCloud\\松霖人脸数据\\face_seg", i+file_name, file_name+".ply")
                 )
    #     break
    # break

